package com.xtayfjpk.zookeeper.test.master_workers.curator;

import java.io.Closeable;
import java.io.IOException;
import java.util.List;
import java.util.Random;
import java.util.concurrent.CountDownLatch;

import org.apache.curator.RetryPolicy;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.recipes.cache.ChildData;
import org.apache.curator.framework.recipes.cache.PathChildrenCache;
import org.apache.curator.framework.recipes.leader.LeaderSelector;
import org.apache.curator.framework.recipes.leader.LeaderSelectorListener;
import org.apache.curator.framework.state.ConnectionState;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.log4j.Logger;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.data.Stat;

import com.xtayfjpk.zookeeper.test.master_workers.curator.listener.ErrorsListener;
import com.xtayfjpk.zookeeper.test.master_workers.curator.listener.MasterListener;
import com.xtayfjpk.zookeeper.test.master_workers.curator.listener.TasksCacheListener;
import com.xtayfjpk.zookeeper.test.master_workers.curator.listener.WorkersCacheListener;
import com.xtayfjpk.zookeeper.test.master_workers.recovery.RecoveredAssignments;
import com.xtayfjpk.zookeeper.test.master_workers.recovery.RecoveryCallback;

public class CuratorMaster implements Closeable, LeaderSelectorListener {
	private static final Logger LOGGER = Logger.getLogger(CuratorMaster.class);
	private static final String hostPort = "192.168.0.202:2181";
	
	private String myId;
	private CuratorFramework client;
	private final LeaderSelector leaderSelector;
	private final PathChildrenCache workersCache;
	private final PathChildrenCache tasksCache;
	
	private CountDownLatch leaderLatch = new CountDownLatch(1);
    private CountDownLatch closeLatch = new CountDownLatch(1);
    private CountDownLatch recoveryLatch = new CountDownLatch(0);
    
    private Random rand = new Random(System.currentTimeMillis());
	
    public CuratorMaster(String myId, String hostPort, RetryPolicy retryPolicy) {
    	LOGGER.info( myId + ": " + hostPort );
    	this.myId = myId;
    	this.client = CuratorFrameworkFactory.newClient(hostPort, retryPolicy);
    	this.leaderSelector = new LeaderSelector(client, "/master", this);
    	this.workersCache = new PathChildrenCache(this.client, "/workers", true);
        this.tasksCache = new PathChildrenCache(this.client, "/tasks", true);
    }
    
    public void startZK() {
    	client.start();
    }
    
    public void bootstrap() throws Exception {
    	Stat stat = client.checkExists().forPath("/workers");
    	if(stat==null) {
    		client.create().forPath("/workers", new byte[0]);
    	}
    	
    	stat = client.checkExists().forPath("/assign");
    	if(stat==null) {
    		client.create().forPath("/assign", new byte[0]);
    	}
    	
    	stat = client.checkExists().forPath("/tasks");
    	if(stat==null) {
    		client.create().forPath("/tasks", new byte[0]);
    	}
    	
    	stat = client.checkExists().forPath("/status");
    	if(stat==null) {
    		client.create().forPath("/status", new byte[0]);
    	}
    }
    
    public void runForMaster() {
    	leaderSelector.setId(myId);
    	LOGGER.info( "Starting master selection: " + myId);
    	//leaderSelector.autoRequeue();
    	leaderSelector.start();
    }
    
    public void awaitLeadership() throws InterruptedException {
    	leaderLatch.await();
    }
    
    public boolean isLeader() {
    	return leaderSelector.hasLeadership();
    }
    
    public void deleteAssignment(String path) throws Exception {
    	LOGGER.info( "Deleting assignment: " + path );
		
    	client.delete().inBackground().forPath(path);
	}
    
    public void assignTasks(List<String> tasks) throws Exception {
    	for(String task : tasks) {
    		assignTask(task, client.getData().forPath("/tasks/" + task));
    	}
	}
    
	public void assignTask(String task, byte[] data) throws Exception {
		List<ChildData> workerList = workersCache.getCurrentData();		
		
		LOGGER.info(String.format("Assigning task %s, data %s", task, new String(data)));
		String designatedWorker = workerList.get(rand.nextInt(workerList.size())).getPath().replaceFirst("/workers/", "");
		
		/*
         * Assign task to randomly chosen worker.
         */
        String path = "/assign/" + designatedWorker + "/" + task;
        createAssignment(path, data);
	}

	public void createAssignment(String path, byte[] data) throws Exception {
		client.create().withMode(CreateMode.PERSISTENT).inBackground().forPath(path, data);	
	}
	
	public void deleteTask(String number) throws Exception {
		LOGGER.info("Deleting task: " + number);
		client.delete().inBackground().forPath("/tasks/task-" + number);
		
		recoveryLatch.countDown();
	}
	
	public void getAbsentWorkerTasks(String absentWorkerPath) throws Exception {
		client.getChildren().inBackground().forPath("/assign/" + absentWorkerPath);		
	}

	@Override
	public void stateChanged(CuratorFramework client, ConnectionState newState) {
		switch (newState) {
		case CONNECTED:
			LOGGER.warn("CONNECTED");
			break;
		case RECONNECTED:
			LOGGER.warn("RECONNECTED");
			break;
		case SUSPENDED:
			LOGGER.warn("Session suspended");
			break;
		case LOST:
			LOGGER.warn("Session LOST");
			try{
                close();
            } catch (IOException e) {
                LOGGER.warn( "Exception while closing", e );
            }
			break;
		case READ_ONLY:
			break;

		default:
			break;
		}
		
	}

	@Override
	public void takeLeadership(CuratorFramework client) throws Exception {
		LOGGER.info( "Mastership participants: " + myId + ", " + leaderSelector.getParticipants() );

		/*
         * Register listeners
         */
		client.getCuratorListenable().addListener(new MasterListener(this));
		client.getUnhandledErrorListenable().addListener(new ErrorsListener(this));
		
		workersCache.getListenable().addListener(new WorkersCacheListener(this));
		workersCache.start();
		
		new RecoveredAssignments(client.getZookeeperClient().getZooKeeper()).recover(new RecoveryCallback() {
			
			@Override
			public void recoveryComplete(int rc, List<String> tasks) {
				try {
					if(rc == RecoveryCallback.FAILED) {
	                    LOGGER.warn("Recovery of assigned tasks failed.");
	                } else {
	                    LOGGER.info( "Assigning recovered tasks" );
	                    recoveryLatch = new CountDownLatch(tasks.size());
	                    assignTasks(tasks);
	                }
					
					new Thread(new Runnable() {
						
						@Override
						public void run() {
							try {
								recoveryLatch.await();
								tasksCache.getListenable().addListener(new TasksCacheListener(CuratorMaster.this));
								tasksCache.start();
							} catch(Exception e) {
								LOGGER.warn("Exception while assigning and getting tasks.", e  );
							}
							
						}
					}).start();
					
					
					leaderLatch.countDown();
				} catch (Exception e) {
					LOGGER.error("Exception while executing the recovery callback", e);
				}
				
			}
		});
		
		/*
         * This latch is to prevent this call from exiting. If we exit, then
         * we release mastership.
         */
		closeLatch.await();
	}

	@Override
	public void close() throws IOException {
		LOGGER.info( "Closing" );
		closeLatch.countDown();
		leaderSelector.close();
		client.close();
	}

	

	
	@SuppressWarnings("resource")
	public static void main(String[] args) throws Exception {
		CuratorMaster master = new CuratorMaster(
				"zhangjin", 
				hostPort, 
				new ExponentialBackoffRetry(1000, 5));
		
		master.startZK();
		master.bootstrap();
		master.runForMaster();
		
		Object object = new Object();
		synchronized (object) {
			object.wait();
		}
	}
	

}
